#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
# Copyright (c) 2021-present, sun-shine cloud, Inc.
# All rights reserved.
# Author：  xiaobo
# Date：    2021-04-21
# rtmp处理类
import rtmplog
import time
import datetime
import public
from dataStat import dstat 

class rtmpStream():

    b_time = 0
    data = b""
    client = None
    app = None          #应用名；
    stream = None       #流名称
    chunksize = 128     #chunk的大小默认是128， 这是默认的大小
    windowsize = 5000000
    live = None
    msg_len = 0
    msg_streamid = 0
    msg_type = 0
    fmt = 0
    csid = 0


    def __init__(self, client, ipaddr):
        t = time.time()
        #秒级时间戳
        ts = int(round(t * 1000))        
        self.b_time = ts
        self.client = client
        self.ipaddr = ipaddr
        return 

    def handle_Rtmp_Stream(self):
        try:
            self.rtmp_handshake()
            ds = None
            while True:
                cmd, msg, astr = self.rev_rtmp_message()
                if not msg and not cmd:
                    break
                if cmd == b"connect":
                    self.send_connect_to_client(msg)
                if cmd == b"createStream":
                    self.send_onFCPublic_to_client()
                if cmd == b"releaseStream":
                    self.stream = astr[1]      #取sttream的名称
                if cmd == b"publish":
                    self.send_public_to_client(msg)   #发送给前置机"_result"信息 
                if cmd == b"chunksize":
                    self.chunksize = msg
                if cmd == b"FCUnpublish":
                    continue
                if cmd == b"video":
                    if not ds:
                        ds = dstat(self.app, self.stream)
                    ds.add_vdat(msg, self.msg_timestamp)
                if cmd == b"audio":
                    if not ds:
                        ds = dstat(self.app, self.stream)
                    ds.add_adat(msg, self.msg_timestamp)
        except Exception as e:
            rtmplog.logger.error(e)
            self.client.close()
            return

    #接收一个message信息；包括 rtmp_head + rtmp_body；
    # RTMP协议定义每个message由remp_head头和rtmp_body构成，每个rtmp_body的长度最大为chunksize构成，但rtmp_head分为四种情况，所以每个message的长度是不一样的；
    # 四种rtmp_head的情况由rtmp_head的第一个字节进行控制（fmt[0-2])
    # 所以在处理包的时候，需要先读取rtmp_head的信息，然后才能够读取rtmp_body的信息； 
    def rev_rtmp_message1(self):

        msg_body = b""
        rm = 0
        while True:
            rtmp_head0 = int.from_bytes(self.rec_msg(1), byteorder='big', signed=False)
            # 一个rtmp的message构成: rtmp_head+rtmp_body
            # rtmp_head：Basic_header+message_header; message_header:Message Type| Payload length|Timestamp|Stream ID
            #读第一个字节，取fmt和csid
            fmt = rtmp_head0 >> 6
            csid = rtmp_head0&0x3F
            #如果csid==0,则csid是后一位的值+64
            if csid == 0:
                csid1 = int.from_bytes(self.rec_msg(1), byteorder='big', signed=False)
                csid = csid1 + 64
            if csid == 1:
                csid1 = int.from_bytes(self.rec_msg(2), byteorder='big', signed=False)
                csid = csid1 + 64
            # rtmplog.logger.info("fmt:"+str(fmt)+"   csid:"+str(csid))
            '''
                fmt=0：长度11 byte，
                在一个块流的开始和时间戳返回的时候必须有这种块
                fmt=1：长度为7 byte，与fmt=0时，相比，该类型少了Message Stream Id，
                具有可变大小消息的流,在第一个消息之后的每个消息的第一个块应该使用这个格式
                fmt=2：长度3 byte，不包含Message Stream Id和Message Length
                具有固定大小消息的流,在第一个消息之后的每个消息的第一个块应该使用这个格式
                fmt=3：长度为0 byte，
                当一个消息被分成多个块,除了第一块以外,所有的块都应使用这种类型
            '''
            self.fmt = fmt
            self.csid = csid

            if fmt == 0:    #会包含4个部分: Timestamp(3),Payload length(3), Message Type(1),Stream ID(4)
                self.msg_timestamp = int.from_bytes(self.rec_msg(3), byteorder='big', signed=False)
                self.msg_len = int.from_bytes(self.rec_msg(3), byteorder='big', signed=False)
                if rm == 0:
                    rm = self.msg_len
                self.msg_type = int.from_bytes(self.rec_msg(1), byteorder='big', signed=False)
                self.msg_streamid = int.from_bytes(self.rec_msg(4), byteorder='little', signed=False)

            elif fmt == 1:    #会包含3个部分: Timestamp(3),Payload length(3), Message Type(1)
                self.msg_timestamp = int.from_bytes(self.rec_msg(3), byteorder='big', signed=False)
                self.msg_len = int.from_bytes(self.rec_msg(3), byteorder='big', signed=False)
                if rm == 0:
                    rm = self.msg_len
                self.msg_type = int.from_bytes(self.rec_msg(1), byteorder='big', signed=False)

            elif fmt == 2:    #会包含1个部分: Timestamp(3)
                self.msg_timestamp = int.from_bytes(self.rec_msg(3), byteorder='big', signed=False)
                if rm == 0:
                    rm = self.msg_len
            else:
                if rm == 0:
                    rm = self.msg_len
            rec = b""
            #下面代码是把一个消息全部的都读完；
            if rm >= self.chunksize:
                rec = self.rec_msg(self.chunksize)
                rm = rm - self.chunksize
            else:
                rec = self.rec_msg(rm)
            msg_body = msg_body + rec
            if len(msg_body) >= self.msg_len:
                break

        # rtmplog.logger.info("body:"+str(self.msg_len)+" msg_body:"+str(len(msg_body)))

        if self.msg_type == 0x11 or self.msg_type == 0x14:   #对于command类型的消息，提取command内容；
            p_str, p_obj =  public.parase_AFM_Data(msg_body)
            return p_str[0], p_obj, p_str
        if self.msg_type == 0x01:   #提取chunksize的设置值
            cksize =  public.parase_ChunkSize_Data(msg_body)
            return b"chunksize", cksize, None
        if self.msg_type == 0x08:   #对于command类型的消息，提取command内容；
            return b"audio", msg_body, None
        if self.msg_type == 0x09:   #对于command类型的消息，提取command内容；
            return b"video", msg_body, None
        return "no_no", None, None


    #接收一个message信息；包括 rtmp_head + rtmp_body；
    # RTMP协议定义每个message由remp_head头和rtmp_body构成，每个rtmp_body的长度最大为chunksize构成，但rtmp_head分为四种情况，所以每个message的长度是不一样的；
    # 四种rtmp_head的情况由rtmp_head的第一个字节进行控制（fmt[0-2])
    # 所以在处理包的时候，需要先读取rtmp_head的信息，然后才能够读取rtmp_body的信息； 
    def rev_rtmp_message(self):

        msg_body = b""
        rm = 0
        while True:
            basic_head = int.from_bytes(self.rec_msg(1), byteorder='big', signed=False)
            # 一个rtmp的message构成: rtmp_head+rtmp_body
            # rtmp_head：Basic_header+message_header; message_header:Message Type| Payload length|Timestamp|Stream ID
            #读第一个字节，取fmt和csid
            fmt = basic_head >> 6
            csid = basic_head&0x3F
            #如果csid==0,则csid是后一位的值+64
            if csid == 0:
                csid1 = int.from_bytes(self.rec_msg(1), byteorder='big', signed=False)
                csid = csid1 + 64
            if csid == 1:
                csid1 = int.from_bytes(self.rec_msg(2), byteorder='big', signed=False)
                csid = csid1 + 64
            # rtmplog.logger.info("fmt:"+str(fmt)+"   csid:"+str(csid))
            '''
                fmt=0：长度11 byte，
                在一个块流的开始和时间戳返回的时候必须有这种块
                fmt=1：长度为7 byte，与fmt=0时，相比，该类型少了Message Stream Id，
                具有可变大小消息的流,在第一个消息之后的每个消息的第一个块应该使用这个格式
                fmt=2：长度3 byte，不包含Message Stream Id和Message Length
                具有固定大小消息的流,在第一个消息之后的每个消息的第一个块应该使用这个格式
                fmt=3：长度为0 byte，
                当一个消息被分成多个块,除了第一块以外,所有的块都应使用这种类型
            '''
            #CSID 0、1、2 被保留， 3 ~ 8 基本都是固定用途，所以 9 ~ 65599 才用于自定义 csid，但一般我们用不到
            # 所以超出3~8都是异常的值，取其基本包(chunksize)即可
            # if fmt == 3 and self.fmt == 1:
            #     self.rec_msg(self.chunksize -1)
            #     continue           

            # if self.fmt == 2 and fmt == 1:
            #     self.rec_msg(self.chunksize -1)
            #     continue
            # if self.fmt == 2 and fmt == 3:
            #     self.rec_msg(self.chunksize -1)
            #     continue

            # if self.fmt == 3 and fmt == 2:
            #     self.rec_msg(self.chunksize -1)
            #     continue
            # if csid > 6:
            #      self.rec_msg(self.chunksize -1)
            #      continue
            # if fmt >=2 and csid != self.csid:
            #     self.rec_msg(self.chunksize -1)
            #     continue
            # if csid == 2 and self.csid == 3:
            #     self.rec_msg(self.chunksize -1)
            #     continue
            
            self.fmt = fmt
            self.csid = csid
            if fmt == 0:    #会包含4个部分: Timestamp(3),Payload length(3), Message Type(1),Stream ID(4)
                self.msg_timestamp = int.from_bytes(self.rec_msg(3), byteorder='big', signed=False)
                self.msg_len = int.from_bytes(self.rec_msg(3), byteorder='big', signed=False)
                if rm == 0:
                    rm = self.msg_len
                self.msg_type = int.from_bytes(self.rec_msg(1), byteorder='big', signed=False)
                self.msg_streamid = int.from_bytes(self.rec_msg(4), byteorder='little', signed=False)

            elif fmt == 1:    #会包含3个部分: Timestamp(3),Payload length(3), Message Type(1)
                self.msg_timestamp = int.from_bytes(self.rec_msg(3), byteorder='big', signed=False)
                self.msg_len = int.from_bytes(self.rec_msg(3), byteorder='big', signed=False)
                if rm == 0:
                    rm = self.msg_len
                self.msg_type = int.from_bytes(self.rec_msg(1), byteorder='big', signed=False)

            elif fmt == 2:    #会包含1个部分: Timestamp(3)
                self.msg_timestamp = int.from_bytes(self.rec_msg(3), byteorder='big', signed=False)
                if rm == 0:
                    rm = self.msg_len
            else:
                if rm == 0:
                    rm = self.msg_len
            rec = b""
            #下面代码是把一个消息全部的都读完；
            if rm >= self.chunksize:
                rec = self.rec_msg(self.chunksize)
                rm = rm - self.chunksize
            else:
                rec = self.rec_msg(rm)
            msg_body = msg_body + rec
            if len(msg_body) >= self.msg_len:
                break

        # rtmplog.logger.info("body:"+str(self.msg_len)+" msg_body:"+str(len(msg_body)))

        if self.msg_type == 0x11 or self.msg_type == 0x14:   #对于command类型的消息，提取command内容；
            p_str, p_obj =  public.parase_AFM_Data(msg_body)
            return p_str[0], p_obj, p_str
        if self.msg_type == 0x01:   #提取chunksize的设置值
            cksize =  public.parase_ChunkSize_Data(msg_body)
            return b"chunksize", cksize, None
        if self.msg_type == 0x08:   #对于command类型的消息，提取command内容；
            return b"audio", msg_body, None
        if self.msg_type == 0x09:   #对于command类型的消息，提取command内容；
            return b"video", msg_body, None
        return "no_no", None, None

    #获取相应的请求信息，并把结果返回给前端
    def send_connect_to_client(self, obj):
        #获取请求是的流信息 app和stream, 这个信息中一定是含有app stream 和 live. type, tcurl三项的内容；该处因为没搞清逻辑，现在采用硬 编码的方式，以后再进行改进；
        #取app
        try:   #
            self.app = obj[b"app"]
        except Exception as e:
            rtmplog.logger.error(e)
            return
        finally:
            rdat =public.mk_windowsize_msg()
            self.send_msg(rdat)
            rdat1 = public.mk_bandwidth_msg()
            self.send_msg(rdat1)
            # rdat =public.mk_chunksize_msg(self.chunksize)
            # self.send_msg(rdat)
            #不论分析的对错，都向client发送成功的信息；
            # r = public.mk_AFM_String_Data("_result")
            # rdat = public.mk_connect_msg(r)
            # self.send_msg(rdat)

    def send_onFCPublic_to_client(self):

            rdat = public.mk_onFCPublish_msg()
            self.send_msg(rdat)
            dat1 = public.mk_result_msg()
            self.send_msg(dat1)

    #接收到public的信息后，获取app和live
    def send_public_to_client(self, cmd):
        #获取请求是的流信息 app和stream, 这个信息中一定是含有app stream 和 live. type, tcurl三项的内容；该处因为没搞清逻辑，现在采用硬 编码的方式，以后再进行改进；
        #取app
        rdat = public.mk_onStatus_msg()
        self.send_msg(rdat)

    #启动握手，如果握手成功，则需要
    def rtmp_handshake(self):

        #from client recvice C0 从客户端接收C0
        c0 = self.rec_msg(1)
        #from client recvice C0
        re1 = self.rec_msg(1536)
        timestamp1 = re1[:4]
        fillzero1 = re1[4:8]
        frandom1 = re1[8:]
        #send S0 to client
        s0 = c0
        #send S0+S1+S2,服务器发送S0+S1+S2
        self.send_msg(s0)
        self.send_msg(re1)
        self.send_msg(re1)
        #rec C2
        #然后客户端再发送C2
        re1 = self.rec_msg(1536)
        return  

    def rec_msg(self, buffer):
        try:
            recv_d = self.client.recv(buffer) #接收一个信息，并指定接收的大小 为1024字节
            #在接收的过程中，发现有些信息会在最后加上c4,暂时不知道怎么处理
            if not recv_d:
                self.client.close()
                exit
            return recv_d
        except Exception as e:
            rtmplog.logger.error("接收信息错误！")
            rtmplog.logger.error(e)
            raise Exception("接收信息错误！")

    def send_msg(self, data):
        try:
            sd = self.client.send(data) #接收一个信息，并指定接收的大小 为1024字节
            # rtmplog.logger.info("发送信息成功，长度是:["+str(len(data))+"]")
            # rtmplog.logger.info(data)
            return True

        except Exception as e:
            rtmplog.logger.error("发送信息错误！")
            rtmplog.logger.error(e)
            raise Exception("发送信息错误！")

    def cur_timestamp(self):

        t = time.time()
        #秒级时间戳
        ts = int(round(t * 1000))
        ts1 = ts - self.b_time
        c_ts = ts1.to_bytes(4, byteorder='big')
        return c_ts